package com.ruoyi.admin.web.controller.tool;

import javax.imageio.*;
import javax.imageio.metadata.IIOMetadata;
import javax.imageio.metadata.IIOMetadataNode;
import javax.imageio.stream.ImageOutputStream;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.awt.image.IndexColorModel;
import java.io.*;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.Stream;

public class ImageCompressor {
    // 支持的图片格式
    private static final List<String> SUPPORTED_FORMATS = Arrays.asList("jpg", "jpeg", "png", "gif");
    // PNG压缩时使用的DPI（降低DPI可以减小文件大小）
    private static final int PNG_TARGET_DPI = 72;
    // PNG压缩尝试的最大颜色数
    private static final int MAX_PNG_COLORS = 256;

    public static void main(String[] args) {
        // 设置源目录和目标目录
        String sourceDir = "D:\\rom-dev\\wallpaper\\common";
        String targetDir = "D:\\rom-dev\\wallpaper\\ok";

        // 设置压缩质量 (0.0f 到 1.0f 之间，值越小压缩率越高)
        float quality = 0.6f;

        try {
            compressImagesInDirectory(sourceDir, targetDir, quality);
            System.out.println("图片压缩完成！");
        } catch (IOException e) {
            System.err.println("压缩过程中发生错误: " + e.getMessage());
            e.printStackTrace();
        }
    }

    /**
     * 压缩指定目录中的所有图片
     */
    public static void compressImagesInDirectory(String sourceDir, String targetDir, float quality) throws IOException {
        Path sourcePath = Paths.get(sourceDir);
        if (!Files.exists(sourcePath) || !Files.isDirectory(sourcePath)) {
            System.err.println("源目录不存在: " + sourceDir);
            return;
        }

        Path targetPath = Paths.get(targetDir);
        Files.createDirectories(targetPath);

        List<File> imageFiles;
        try (Stream<Path> walk = Files.walk(sourcePath)) {
            imageFiles = walk
                    .filter(Files::isRegularFile)
                    .map(Path::toFile)
                    .filter(file -> isSupportedImageFormat(file.getName()))
                    .collect(Collectors.toList());
        }

        for (File sourceFile : imageFiles) {
            try {
                long originalSize = sourceFile.length();
                File targetFile = targetPath.resolve(sourceFile.getName()).toFile();

                compressImage(sourceFile, targetFile, quality);

                long compressedSize = targetFile.length();
                if (compressedSize > originalSize) {
                    Files.copy(sourceFile.toPath(), targetFile.toPath(),
                            java.nio.file.StandardCopyOption.REPLACE_EXISTING);
                    System.out.printf("使用原始文件(压缩后更大): %s%n", sourceFile.getName());
                } else {
                    System.out.printf("已压缩: %s -> %s (%.2f%% 原始大小)%n",
                            sourceFile.getAbsolutePath(),
                            targetFile.getAbsolutePath(),
                            (double)compressedSize / originalSize * 100);
                }
            } catch (Exception e) {
                System.err.printf("压缩失败: %s - %s%n",
                        sourceFile.getAbsolutePath(), e.getMessage());
            }
        }
    }

    /**
     * 压缩单个图片文件
     */
    public static void compressImage(File sourceFile, File targetFile, float quality) throws IOException {
        quality = Math.max(0.0f, Math.min(1.0f, quality));
        String formatName = getFileExtension(sourceFile.getName()).toLowerCase();

        if ("gif".equalsIgnoreCase(formatName)) {
            Files.copy(sourceFile.toPath(), targetFile.toPath(),
                    java.nio.file.StandardCopyOption.REPLACE_EXISTING);
            return;
        }

        BufferedImage image = ImageIO.read(sourceFile);

        if ("jpg".equalsIgnoreCase(formatName) || "jpeg".equalsIgnoreCase(formatName)) {
            compressJPEG(image, targetFile, quality);
        } else if ("png".equalsIgnoreCase(formatName)) {
            compressPNG(image, targetFile, quality);
        } else {
            compressDefault(image, targetFile, formatName, quality);
        }
    }

    /**
     * 压缩JPEG图片
     */
    private static void compressJPEG(BufferedImage image, File targetFile, float quality) throws IOException {
        Iterator<ImageWriter> writers = ImageIO.getImageWritersByFormatName("JPEG");
        if (!writers.hasNext()) {
            throw new IllegalStateException("找不到JPEG写入器");
        }
        ImageWriter writer = writers.next();

        try (OutputStream os = new FileOutputStream(targetFile);
             ImageOutputStream ios = ImageIO.createImageOutputStream(os)) {

            writer.setOutput(ios);

            ImageWriteParam param = writer.getDefaultWriteParam();
            param.setCompressionMode(ImageWriteParam.MODE_EXPLICIT);
            param.setCompressionQuality(quality);

            ImageTypeSpecifier typeSpecifier = ImageTypeSpecifier.createFromRenderedImage(image);
            IIOMetadata metadata = writer.getDefaultImageMetadata(typeSpecifier, param);

            writer.write(metadata, new IIOImage(image, null, metadata), param);
        } finally {
            writer.dispose();
        }
    }

    /**
     * 压缩PNG图片 - 优化版
     */
    private static void compressPNG(BufferedImage image, File targetFile, float quality) throws IOException {
        // 计算压缩级别 (0-9)
        int compressionLevel = (int) ((1.0f - quality) * 9);

        // 尝试不同的压缩策略并选择最佳结果
        File tempFile = File.createTempFile("png_compress_", ".png");
        tempFile.deleteOnExit();

        // 策略1: 尝试使用索引色模式 (8位)
        BufferedImage indexedImage = convertToIndexedImage(image, MAX_PNG_COLORS);
        boolean indexedSuccess = false;

        if (indexedImage != null) {
            try {
                savePNGWithCompression(indexedImage, tempFile, compressionLevel);
                if (tempFile.length() < targetFile.length() || !targetFile.exists()) {
                    Files.copy(tempFile.toPath(), targetFile.toPath(),
                            java.nio.file.StandardCopyOption.REPLACE_EXISTING);
                    indexedSuccess = true;
                }
            } catch (Exception e) {
                System.err.println("索引色PNG保存失败: " + e.getMessage());
            }
        }

        // 策略2: 如果索引色模式不成功或不适用，尝试真彩色模式
        if (!indexedSuccess) {
            savePNGWithCompression(image, targetFile, compressionLevel);
        }
    }

    /**
     * 使用指定压缩级别保存PNG图片
     */
    private static void savePNGWithCompression(BufferedImage image, File targetFile, int compressionLevel) throws IOException {
        Iterator<ImageWriter> writers = ImageIO.getImageWritersByFormatName("PNG");
        if (!writers.hasNext()) {
            throw new IllegalStateException("找不到PNG写入器");
        }
        ImageWriter writer = writers.next();

        try (OutputStream os = new FileOutputStream(targetFile);
             ImageOutputStream ios = ImageIO.createImageOutputStream(os)) {

            writer.setOutput(ios);

            ImageWriteParam param = writer.getDefaultWriteParam();

            ImageTypeSpecifier typeSpecifier = ImageTypeSpecifier.createFromRenderedImage(image);
            IIOMetadata metadata = writer.getDefaultImageMetadata(typeSpecifier, param);

            // 设置压缩参数
            if (metadata.isReadOnly()) {
                metadata = writer.getDefaultImageMetadata(typeSpecifier, null);
            }

            // 设置DPI
            setPNGDPI(metadata);

            // 移除不必要的元数据
            removeMetadata(metadata);

            writer.write(metadata, new IIOImage(image, null, metadata), param);
        } finally {
            writer.dispose();
        }
    }

    /**
     * 将图像转换为索引色模式
     */
    private static BufferedImage convertToIndexedImage(BufferedImage image, int maxColors) {
        // 分析图像中的颜色数量
        Set<Integer> colorSet = new HashSet<>();
        int width = image.getWidth();
        int height = image.getHeight();

        // 限制采样点数量，避免处理过大图像时性能问题
        int sampleStep = Math.max(1, (width * height) / 10000);

        for (int y = 0; y < height; y += sampleStep) {
            for (int x = 0; x < width; x += sampleStep) {
                colorSet.add(image.getRGB(x, y));
                if (colorSet.size() > maxColors) {
                    // 颜色太多，不适合转换为索引色
                    return null;
                }
            }
        }

        // 如果颜色数量少，可以转换为索引色
        if (colorSet.size() <= maxColors) {
            byte[] r = new byte[colorSet.size()];
            byte[] g = new byte[colorSet.size()];
            byte[] b = new byte[colorSet.size()];

            int index = 0;
            for (int rgb : colorSet) {
                r[index] = (byte) ((rgb >> 16) & 0xFF);
                g[index] = (byte) ((rgb >> 8) & 0xFF);
                b[index] = (byte) (rgb & 0xFF);
                index++;
            }

            IndexColorModel colorModel = new IndexColorModel(8, colorSet.size(), r, g, b);
            BufferedImage indexedImage = new BufferedImage(width, height, BufferedImage.TYPE_BYTE_INDEXED, colorModel);

            Graphics2D g2d = indexedImage.createGraphics();
            g2d.drawImage(image, 0, 0, null);
            g2d.dispose();

            return indexedImage;
        }

        return null;
    }

    /**
     * 设置PNG图片的DPI
     */
    private static void setPNGDPI(IIOMetadata metadata) throws IIOException {
        if (!metadata.isStandardMetadataFormatSupported()) {
            return;
        }

        IIOMetadataNode root = new IIOMetadataNode("javax_imageio_1.0");
        IIOMetadataNode dimension = new IIOMetadataNode("Dimension");
        IIOMetadataNode horizontalPixelSize = new IIOMetadataNode("HorizontalPixelSize");
        IIOMetadataNode verticalPixelSize = new IIOMetadataNode("VerticalPixelSize");

        double pixelSize = 25.4 / PNG_TARGET_DPI;
        horizontalPixelSize.setAttribute("value", Double.toString(pixelSize));
        verticalPixelSize.setAttribute("value", Double.toString(pixelSize));

        dimension.appendChild(horizontalPixelSize);
        dimension.appendChild(verticalPixelSize);
        root.appendChild(dimension);

        metadata.mergeTree("javax_imageio_1.0", root);
    }

    /**
     * 移除不必要的元数据
     */
    private static void removeMetadata(IIOMetadata metadata) {
        if (!metadata.isStandardMetadataFormatSupported()) {
            return;
        }

        try {
            // 创建一个空的元数据树
            IIOMetadataNode root = new IIOMetadataNode("javax_imageio_1.0");
            metadata.mergeTree("javax_imageio_1.0", root);
        } catch (IIOException e) {
            System.err.println("移除元数据失败: " + e.getMessage());
        }
    }

    /**
     * 默认压缩方法
     */
    private static void compressDefault(BufferedImage image, File targetFile, String formatName, float quality) throws IOException {
        Iterator<ImageWriter> writers = ImageIO.getImageWritersByFormatName(formatName);
        if (!writers.hasNext()) {
            throw new IllegalStateException("找不到支持的图片写入器: " + formatName);
        }
        ImageWriter writer = writers.next();

        try (OutputStream os = new FileOutputStream(targetFile);
             ImageOutputStream ios = ImageIO.createImageOutputStream(os)) {

            writer.setOutput(ios);

            ImageWriteParam param = writer.getDefaultWriteParam();
            if (param.canWriteCompressed()) {
                param.setCompressionMode(ImageWriteParam.MODE_EXPLICIT);
                param.setCompressionQuality(quality);
            }

            writer.write(null, new IIOImage(image, null, null), param);
        } finally {
            writer.dispose();
        }
    }

    /**
     * 检查文件是否为支持的图片格式
     */
    private static boolean isSupportedImageFormat(String fileName) {
        String extension = getFileExtension(fileName).toLowerCase();
        return SUPPORTED_FORMATS.contains(extension);
    }

    /**
     * 获取文件扩展名
     */
    private static String getFileExtension(String fileName) {
        int dotIndex = fileName.lastIndexOf('.');
        return (dotIndex == -1) ? "" : fileName.substring(dotIndex + 1);
    }
}